package com.example.dashscopetest.controller;

import com.alibaba.cloud.ai.dashscope.agent.DashScopeAgent;
import com.alibaba.cloud.ai.dashscope.agent.DashScopeAgentOptions;
import com.alibaba.cloud.ai.dashscope.api.DashScopeAgentApi;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@RestController
@RequestMapping("ai")
public class BailianAgentRagStreamController {

    private static final Logger LOGGER = LoggerFactory.getLogger(BailianAgentRagStreamController.class);

    @Value("${spring.ai.dashscope.agent.app-id}")
    private String appId;

    private DashScopeAgent agent;

    public BailianAgentRagStreamController(DashScopeAgentApi dashScopeAgentApi) {
        this.agent = new DashScopeAgent(dashScopeAgentApi, DashScopeAgentOptions.builder().withSessionId("current_session_id").withIncrementalOutput(true).withHasThoughts(true).build());
    }

    @GetMapping(value = "/bailian/agent/stream", produces = "text/event-stream")
    public Flux<String> streamRag(@RequestParam(value = "message", defaultValue = "你好，请问你的知识库文档主要是关于什么内容的?") String message) {
        return agent.stream(new Prompt(message, DashScopeAgentOptions.builder().withAppId(appId).build())).map(response -> {
            if (response == null || response.getResult() == null) {
                LOGGER.error("chat response is null");
                return "chat response is null";
            }
            AssistantMessage output = response.getResult().getOutput();
            return output.getContent();
        });
    }

}
